/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include <sstream>
#include <cstring>
#include "utils.h"

namespace atb_plugin {

atb::SVector<int64_t> GetCopyTensorStride(atb::Dims &tensorDims)
{
    atb::SVector<int64_t> tmpStrides(tensorDims.dimNum, 1);
    for (int64_t i = static_cast<int64_t>(tensorDims.dimNum) - 2; i >= 0; i--) {
        tmpStrides[i] = tensorDims.dims[i + 1] * tmpStrides[i + 1];
    }
    return tmpStrides;
}

atb::Tensor SqueezeBatchSeq(atb::Tensor atbTensor)
{
    if (atbTensor.desc.shape.dimNum == DIM3) {
        atbTensor.desc.shape.dimNum = DIM2;
        atbTensor.desc.shape.dims[DIM0] = atbTensor.desc.shape.dims[DIM0] * atbTensor.desc.shape.dims[DIM1];
        atbTensor.desc.shape.dims[DIM1] = atbTensor.desc.shape.dims[DIM2];
    }
    return atbTensor;
}

} // namespace atb_plugin
